from function import *
from decimal import Decimal, ROUND_HALF_UP

def cal_macd(df):
    cond = df['股票名称'].str.contains('ST')
    df['涨停价'] = df['前收盘价'] * 1.1
    df['跌停价'] = df['前收盘价'] * 0.9
    df.loc[cond, '涨停价'] = df['前收盘价'] * 1.05
    df.loc[cond, '跌停价'] = df['前收盘价'] * 0.95

    # 2020年8月3日之后涨跌停规则有所改变
    # 新规的科创板
    new_rule_kcb = (df['交易日期'] > pd.to_datetime('2020-08-03')) & df['股票代码'].str.contains('sh68')
    # 新规的创业板
    new_rule_cyb = (df['交易日期'] > pd.to_datetime('2020-08-03')) & df['股票代码'].str.contains('sz30')
    # 北交所条件
    cond_bj = df['股票代码'].str.contains('bj')

    # 科创板 & 创业板
    df.loc[new_rule_kcb | new_rule_cyb, '涨停价'] = df['前收盘价'] * 1.2
    df.loc[new_rule_kcb | new_rule_cyb, '跌停价'] = df['前收盘价'] * 0.8

    # 北交所
    df.loc[cond_bj, '涨停价'] = df['前收盘价'] * 1.3
    df.loc[cond_bj, '跌停价'] = df['前收盘价'] * 0.7

    # 四舍五入
    df['涨停价'] = df['涨停价'].apply(
        lambda x: float(Decimal(x * 100).quantize(Decimal('1'), rounding=ROUND_HALF_UP) / 100))
    df['跌停价'] = df['跌停价'].apply(
        lambda x: float(Decimal(x * 100).quantize(Decimal('1'), rounding=ROUND_HALF_UP) / 100))

    # 判断是否一字涨停
    df['一字涨停'] = False
    df.loc[df['最低价'] >= df['涨停价'], '一字涨停'] = True
    # 判断是否一字跌停
    df['一字跌停'] = False
    df.loc[df['最高价'] <= df['跌停价'], '一字跌停'] = True
    # 判断是否开盘涨停
    df['开盘涨停'] = False
    df.loc[df['开盘价'] >= df['涨停价'], '开盘涨停'] = True
    # 判断是否开盘跌停
    df['开盘跌停'] = False
    df.loc[df['开盘价'] <= df['跌停价'], '开盘跌停'] = True

    # ===计算复权价===
    df = rehabilitation(df)

    # ===MACD金叉死叉===
    df['EMA12'] = df['收盘价'].ewm(span=12, adjust=False).mean()
    df['EMA26'] = df['收盘价'].ewm(span=26, adjust=False).mean()
    # 计算DIF
    df['DIF'] = df['EMA12'] - df['EMA26']
    # 计算DEA
    df['DEA'] = df['DIF'].ewm(span=9, adjust=False).mean()
    # 计算MACD柱
    df['MACD'] = 2 * (df['DIF'] - df['DEA'])
    con1 = (df['DIF'] > df['DEA']) & (df['DIF'].shift(1) <= df['DEA'].shift(1))
    con2 = (df['DIF'] < df['DEA']) & (df['DIF'].shift(1) >= df['DEA'].shift(1))
    df.loc[con1, 'signal'] = 1
    df.loc[con2, 'signal'] = 0

    # ===MACD水上金叉水下死叉===
    # df['EMA12'] = df['收盘价'].ewm(span=12, adjust=False).mean()
    # df['EMA26'] = df['收盘价'].ewm(span=26, adjust=False).mean()
    # # 计算DIF
    # df['DIF'] = df['EMA12'] - df['EMA26']
    # # 计算DEA
    # df['DEA'] = df['DIF'].ewm(span=9, adjust=False).mean()
    # # 计算MACD柱
    # df['MACD'] = 2 * (df['DIF'] - df['DEA'])
    # con1 = (df['DIF'] > df['DEA']) &( df['DIF'].shift(1) <= df['DEA'].shift(1))
    # con2 = (df['DIF'] < df['DEA']) &( df['DIF'].shift(1) >= df['DEA'].shift(1))
    # con3 = df['DIF'] > 0
    # con4 = df['DEA'] > 0
    # con5 = df['DIF'] < 0
    # con6 = df['DEA'] < 0
    # df.loc[con1 & con3 & con4, 'signal'] = 1
    # df.loc[con2 & con5 & con6, 'signal'] = 0

    # # ===新周期MACD及黄柱计算===
    # df['EMA14'] = df['收盘价'].ewm(span=14, adjust=False).mean()
    # df['EMA53'] = df['收盘价'].ewm(span=53, adjust=False).mean()
    # # 计算DIF
    # df['DIF'] = df['EMA14'] - df['EMA53']
    # # 计算DEA
    # df['DEA'] = df['DIF'].ewm(span=5, adjust=False).mean()
    # # 计算MACD柱
    # df['MACD'] = 2 * (df['DIF'] - df['DEA'])
    # df['MACD_ma5'] = df['MACD'].rolling(5).mean()
    # con1 = df['MACD'] > df['MACD_ma5']
    # con2 = df['MACD'].shift(1) <= df['MACD_ma5'].shift(1)
    # df.loc[con1 & con2, 'signal'] = 1

    # # ===MACD顶底背离===
    # df['EMA12'] = df['收盘价'].ewm(span=12, adjust=False).mean()
    # df['EMA26'] = df['收盘价'].ewm(span=26, adjust=False).mean()
    # # 计算DIF
    # df['DIF'] = df['EMA12'] - df['EMA26']
    # # 计算DEA
    # df['DEA'] = df['DIF'].ewm(span=9, adjust=False).mean()
    # # 计算MACD柱
    # df['MACD'] = 2 * (df['DIF'] - df['DEA'])
    # df['max'] = df[['开盘价', '收盘价']].max(axis=1)
    # # 计算开盘和收盘的最低价
    # df['min'] = df[['开盘价', '收盘价']].min(axis=1)
    #
    # # 峰值条件：max大于前后两天，且max大于最近30天的所有max
    # df.loc[(df['max'].shift(1) > df['max']) & (df['max'].shift(1) > df['max'].shift(2)) & (
    #         df['max'].shift(1) == df['max'].rolling(30).max()), 'price_new_high'] = 1
    # # 谷值条件：min小于前后两天，且小于最近30天所有的min
    # df.loc[(df['min'].shift(1) < df['min']) & (df['min'].shift(1) < df['min'].shift(2)) & (
    #         df['min'].shift(1) == df['min'].rolling(30).min()), 'price_new_low'] = 1
    #
    # # 计算DIF高低点
    # df.loc[df['DIF'] == df['DIF'].rolling(30).max(), 'DIF_new_high'] = 1
    # df.loc[df['DIF'] == df['DIF'].rolling(30).min(), 'DIF_new_low'] = 1
    #
    # # 记录前后的高低点
    # df.loc[df['price_new_high'] == 1, 'last_peak_price'] = df['收盘价']
    # df.loc[df['price_new_high'] == 1, 'last_peak_dif'] = df['DIF']
    #
    # df.loc[df['price_new_low'] == 1, 'last_valley_price'] = df['收盘价']
    # df.loc[df['price_new_low'] == 1, 'last_valley_dif'] = df['DIF']
    #
    # # 2.3、填充空值 & 取前值
    # df['last_peak_price'].fillna(method='ffill', inplace=True)
    # df['last_peak_dif'].fillna(method='ffill', inplace=True)
    # df['last_peak_price'] = df['last_peak_price'].shift(1)
    # df['last_peak_dif'] = df['last_peak_dif'].shift(1)
    #
    # df['last_valley_price'].fillna(method='ffill', inplace=True)
    # df['last_valley_dif'].fillna(method='ffill', inplace=True)
    # df['last_valley_price'] = df['last_valley_price'].shift(1)
    # df['last_valley_dif'] = df['last_valley_dif'].shift(1)
    #
    # # 3、计算顶底背离
    # cond1 = df['price_new_high'] == 1  # 条件1：在拐点处判断
    # cond2 = df['price_new_low'] == 1
    # cond3 = df['收盘价'] > df['last_peak_price']  # 条件2：当前均值比上一次拐点高（股价创新高）
    # cond4 = df['收盘价'] < df['last_valley_price']  # 条件3：当前均值比上一次拐点低（股价创新低）
    # cond5 = df['DIF'] > df['last_valley_dif']  # 条件4：当前DEA比上一次拐点高（DEA创新高）
    # cond6 = df['DIF'] < df['last_peak_dif']  # 条件4：当前DEA比上一次拐点低（DEA创新低）
    #
    # df.loc[cond1 & cond3 & cond6, 'state'] = '顶背离'
    # df.loc[cond2 & cond4 & cond5, 'state'] = '底背离'
    #
    # df.loc[cond1 & cond3 & cond6, 'signal'] = 0  # 卖出信号（顶背离）
    # df.loc[cond2 & cond4 & cond5, 'signal'] = 1  # 买入信号（底背离）

    return df


def cal_rsi():
    """RSI指标计算 应该独立出去或形成一个策略攻城列表"""
    pass


def cal_kdj(df, n=9, m=3):
    """KDJ指标计算"""
    pass


